import asyncio
import ssl
import time
import threading
import traceback
from aioquic.asyncio import connect
from aioquic.asyncio.protocol import QuicConnectionProtocol
from aioquic.quic.configuration import QuicConfiguration
from aioquic.h3.connection import H3Connection
from aioquic.h3.events import (
    HeadersReceived,
    WebTransportStreamDataReceived,
    DatagramReceived,
)
from aioquic.quic.events import StreamDataReceived
from collections import deque
from ..utils.logger import get_logger
from ..eztrade_client import PyMessage, Authority

logger = get_logger("eztrade.wt", level="warning")

# %% WT 连接客户端
class WebTransportClient:
    __slots__ = (
        "connection",
        "http3",
        "session_id",
        "connected",
        "_wt_stream_ids",
        "server",
        "port",
        "path",
        "buffer",
        "start_time",
        "stop_flag",
        "bg_thread",
        "bg_loop",
        )

    def __init__(self, buffer_size=65536):
        self.connection = None
        self.http3 = None
        self.session_id = None
        self.connected = False
        self._wt_stream_ids: deque[int] = deque(maxlen=100) # 限制监听的发起流
        self.stop_flag = asyncio.Event()  # 用于停止客户端
        self.buffer:deque[bytes] = deque(maxlen=buffer_size)  # 用于存储接收的数据

    def _quic_event_received(self, event):
        """处理 QUIC 事件"""
        logger.trace(f"收到 QUIC 事件: {type(event).__name__}")

        if isinstance(event, StreamDataReceived):
            if event.stream_id in self._wt_stream_ids:
                # 处理 WebTransport 流数据
                data = event.data
                if data:
                    logger.debug(f"QUIC WebTransport 流 {event.stream_id} 收到数据: {event.data}")

        if self.http3:
            for http_event in self.http3.handle_event(event):
                if isinstance(http_event, HeadersReceived):
                    self._handle_headers_received(http_event)
                elif isinstance(http_event, WebTransportStreamDataReceived):
                    data = http_event.data
                    if data == b"ping":
                        new_stream_id = self.create_webtransport_stream(is_unidirectional=False)
                        self.start_time = time.time() # 刷新时间
                        # 新建双向流中回应
                        self.send_stream_data(new_stream_id, b"pong", end_stream=True)
                    elif data: # 只处理非空数据
                        self.buffer.append(data) # 存储接收的字节流
                elif isinstance(http_event, DatagramReceived):
                    # FIXME: 暂时好像是unreadable的
                    logger.debug(f"收到数据报: {http_event.data.decode()}")

    def _handle_headers_received(self, event: HeadersReceived):
        """处理 HTTP/3 头部"""
        headers = dict(event.headers)
        status = headers.get(b":status")
        logger.debug(f"收到 HTTP/3 头部，状态: {status}")
        if status == b"200":
            self.connected = True
            logger.debug("WebTransport 连接建立成功!")
        elif status in {b"403", b"404", b"421", b"500"}:
            logger.error(f"WebTransport 连接失败，状态码: {status.decode()}")
            self.stop_flag.set()

    async def connect(self, server: str, port: int, path: str, user_id: int = None):
        """建立 WebTransport 连接"""
        config = QuicConfiguration(is_client=True, alpn_protocols=["h3"])
        config.verify_mode = ssl.CERT_NONE
        config.server_name = server

        self.server = server
        self.port = port
        self.path = path
        logger.info(f"正在连接到 WebTransport 服务器 {server}:{port}{path} ...")

        async with connect(
            server, port, configuration=config, wait_connected=True
        ) as protocol:
            self.connection:QuicConnectionProtocol = protocol
            self.http3 = H3Connection(protocol._quic, enable_webtransport=True)
            
            # 设置事件处理器
            protocol.quic_event_received = self._quic_event_received
            
            # 发起 WebTransport CONNECT 请求
            await self._send_webtransport_request()
            
            # 等待连接建立
            self.start_time = time.time()  # 记录起始时间
            logger.debug("等待 WebTransport 会话建立...")
            while not self.connected:
                protocol.transmit()
                await asyncio.sleep(0.1)
                if time.time() - self.start_time > 10:
                    raise TimeoutError("WebTransport 会话超时, 10秒未建立连接")
            
            # 开始持续发送和监听
            await self._run_client_loop(user_id)

    async def _send_webtransport_request(self):
        """发送 WebTransport CONNECT 请求"""
        self.session_id = self.connection._quic.get_next_available_stream_id(is_unidirectional=False)
        
        self.http3.send_headers(
            self.session_id,
            [
                (b":method", b"CONNECT"),
                (b":protocol", b"webtransport"),
                (b":scheme", b"https"),
                (b":authority", f"{self.server}:{self.port}".encode()),
                (b":path", self.path.encode()),
            ],
            end_stream=False,
        )
        
        self.connection.transmit()
        logger.debug(f"发送了 WebTransport CONNECT 请求, 会话ID: {self.session_id}")

    def create_webtransport_stream(self, is_unidirectional=False):
        """创建 WebTransport 流"""
        if not self.connected:
            raise Exception("WebTransport 会话未建立")
        
        stream_id = self.http3.create_webtransport_stream(self.session_id, is_unidirectional)
        self.connection.transmit()
        self._wt_stream_ids.append(stream_id)
        if not is_unidirectional:
            self._wt_stream_ids.append(stream_id + 1)  # 双向流会有两个ID
        return stream_id

    def send_stream_data(self, stream_id, data, end_stream=False):
        """发送流数据"""
        self.http3._quic.send_stream_data(stream_id, data, end_stream=end_stream)
        self.connection.transmit()

    def send_datagram(self, data):
        """发送数据报"""
        self.http3.send_datagram(self.session_id, data)
        self.connection.transmit()

    async def _run_client_loop(self, user_id: int = None):
        """客户端主循环：持续发送和监听"""
        # 创建初始流并发送数据
        stream_id = self.create_webtransport_stream(is_unidirectional=True)
        # self.send_stream_data(stream_id, b"Hello from initial stream!", end_stream=True)
        logger.debug(f"发送了初始流数据, 流ID: {stream_id}")
        
        # 发送初始数据报 Deprecated(不用datagram交互)
        # self.send_datagram(b"Hello from datagram!")
        # logger.debug("发送了初始数据报")

        # 发送鉴权消息(空id)
        self.send_message(PyMessage(Authority(user_id)))

        self.start_time = time.time()  # 记录起始时间
        # 持续循环
        try:
            while True:
                # 持续监听和传输
                self.connection.transmit()
                await asyncio.sleep(5)
                if time.time() - self.start_time > 30:
                    raise TimeoutError("WebTransport 会话超时, 30秒未收到数据")
        except KeyboardInterrupt:
            logger.warning("用户中断，退出...")
        except Exception as e:
            logger.error(f"发生错误: {e}")
        finally:
            self.stop_flag.set()

    def send_message(self, message: PyMessage):
        """发送消息"""
        if not self.connected:
            raise Exception("WebTransport 会话未建立")
        if self.stop_flag.is_set():
            logger.error("WebTransport 会话已停止，无法发送消息")
            return
        
        try:
            stream_id = self.create_webtransport_stream(is_unidirectional=True)
            self.send_stream_data(stream_id, message.to_binary_bytes(), end_stream=True)
            logger.info(f"发送消息: {message}, 流ID: {stream_id}")
        except Exception as e:
            logger.error(f"发送消息失败: {''.join(traceback.format_exception(e))}")
            raise e

    def create(self, server: str, port: int, path: str, user_id: int = None):
        """创建 WebTransport 连接线程"""
        # 启动PorterLoop Thread闭包
        logger.info("创建 WebTransport 连接线程")
        loop_future = asyncio.Future()
        event = threading.Event()
        def start_loop(loop_future: asyncio.Future):
            loop = asyncio.new_event_loop() # 新建线程内的事件循环
            loop_future.set_result(loop)    # 传递事件循环到主线程
            event.set()  # 设置事件标志，表示线程已启动
            loop.run_forever()
        thread = threading.Thread(target=start_loop, args=(loop_future,))
        thread.start()
        event.wait(timeout=1.0)  # 等待线程启动
        loop:asyncio.AbstractEventLoop = loop_future.result() # 获取子线程事件循环
        coroutine = self.connect(server, port, path, user_id)
        asyncio.run_coroutine_threadsafe(coroutine, loop) # 将协程任务提交到子线程事件循环
        self.bg_thread = thread
        self.bg_loop = loop

    def stop_loop_thread(self):
        logger.info("停止 WebTransport 客户端线程...")
        tasks = asyncio.all_tasks(self.bg_loop)
        for t in tasks:
            t.cancel()
        self.bg_loop.stop()
        self.bg_thread.join()